package com.uxsino.commons.tftp;

import java.io.IOException;
import java.io.InputStream;
import java.io.OutputStream;
import java.net.InetAddress;
import java.net.SocketTimeoutException;
import java.util.List;
import java.util.concurrent.ConcurrentHashMap;
import java.util.function.Consumer;
import java.util.function.Function;

import org.apache.commons.net.tftp.TFTP;
import org.apache.commons.net.tftp.TFTPPacket;
import org.apache.commons.net.tftp.TFTPPacketException;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import com.google.common.base.Strings;
import com.google.common.collect.Lists;

/**
 * 
 * 描述：TFTP server
 * @author <a href="mailto:royrxc@gmail.com">Ran</a>
 *
 * @date 2018年3月8日
 */
public class TFTPServer extends Thread {
    public enum ServerMode {
                            GET_ONLY,
                            PUT_ONLY,
                            GET_AND_PUT;
    }

    TFTP server = new TFTP();

    private static final Logger LOGGER = LoggerFactory.getLogger(TFTPServer.class);

    public static final int BUF_SIZE = 512;

    private ConcurrentHashMap<String, TFTPSession> sessions = new ConcurrentHashMap<>();

    private ServerMode mode = ServerMode.GET_AND_PUT;// 默认读写方式

    private int port = TFTP.DEFAULT_PORT;

    private int timeout = TFTP.DEFAULT_TIMEOUT;

    private String host;

    private Function<String, InputStream> readerFn;// 读文件

    private Function<String, OutputStream> writerFn;// 写文件

    private Function<String, List<String>> filesFn;// 根据传递的正则表达式,找目录中的所有文件列表，如果传递的是空，则返回所有的文件路径

    private Consumer<Exception> error;

    private boolean running = true;

    public static TFTPServer of() {
        return new TFTPServer();
    }

    private TFTPServer() {
        super("TFTP-SERVER");
    }

    public TFTPServer error(Consumer<Exception> error) {
        this.error = error;
        return this;
    }

    public TFTPServer mode(ServerMode mode) {
        this.mode = mode;
        return this;
    }

    /**
     * 是否設置為守護進程
     * @param daemon
     * @return
     */
    public TFTPServer deamon(boolean daemon) {
        this.setDaemon(daemon);
        return this;
    }

    public TFTPServer host(String host) {
        this.host = host;
        return this;
    }

    /**
     * 设置服务器的服务端口
     * @param port
     * @return
     */
    public TFTPServer port(int port) {
        if (port > 0) {
            this.port = port;
        }
        return this;
    }

    public TFTPServer readerFn(Function<String, InputStream> readerFn) {
        this.readerFn = readerFn;
        return this;
    }

    public TFTPServer writerFn(Function<String, OutputStream> writerFn) {
        this.writerFn = writerFn;
        return this;
    }

    public TFTPServer timeout(int timeoutMillisecond) {
        this.timeout = timeoutMillisecond;
        return this;
    }

    public TFTPServer filesFn(Function<String, List<String>> filesFn) {
        this.filesFn = filesFn;
        return this;
    }

    public List<String> files(String regExp) {
        if (this.filesFn == null) {
            return Lists.newArrayList();
        }
        return this.filesFn.apply(regExp);
    }

    public void close() {
        this.running = false;
        if (this.server != null) {
            if (this.server.isOpen()) {
                this.server.close();
            }
        }
    }

    public void open() {
        if (this.server == null) {
            this.server = new TFTP();
        }
        try {
            if (!this.server.isOpen()) {
                InetAddress h = null;
                if (!Strings.isNullOrEmpty(this.host)) {
                    try {
                        h = InetAddress.getByName(this.host);
                    } catch (Exception e) {
                    }
                }

                server.open(this.port, h);
            }
            server.setSoTimeout(this.timeout);// 设置不超时
            server.setDefaultTimeout(this.timeout);
        } catch (Exception e) {
            if (this.error != null) {
                this.error.accept(e);
            }
            LOGGER.error("Open tftp server error: ", e);
        }
    }

    public void closeOtherSamePort() {
        Thread.getAllStackTraces().forEach((t, ees) -> {
            if (t.getName().equals(this.getName()) && t instanceof TFTPServer && t != this) {
                ((TFTPServer) t).close();
            }
        });
    }

    @Override
    public void run() {
        this.open();
        while (this.running) {
            if (server.isOpen()) {
                try {
                    TFTPPacket packet = this.server.receive();
                    LOGGER.debug("Tftp server received: " + packet.toString());
                    this.process(packet);
                } catch (SocketTimeoutException e) {
                    LOGGER.debug("Tftp server timeout, waiting...");
                } catch (IOException | TFTPPacketException e) {
                    LOGGER.error("Tftp server error: ", e);
                }
            } else {
                try {
                    Thread.sleep(5000);
                    this.open();
                } catch (InterruptedException e) {
                    LOGGER.error("tftp thread sleep error: ", e);
                }
            }
        }
        this.close();
    }

    public void process(TFTPPacket packet) throws TFTPPacketException {
        String key = packet.getAddress().getHostAddress() + ":" + packet.getPort();
        LOGGER.debug("session count : " + this.sessions.size() + " Tftp connected : " + key);

        TFTPSession session = this.sessions.get(key);
        if (session == null) {
            session = TFTPSession.of().mode(mode).readerFn(readerFn).writerFn(writerFn).sender(p -> {
                try {
                    this.server.send(p);
                } catch (IOException e) {
                    LOGGER.error("send data error: ", e);
                }
            });
        }
        if (session.process(packet)) {
            if (session.isEnd()) {
                this.sessions.remove(key);
            } else {
                this.sessions.put(key, session);
            }
        } else {
            if (session.isEnd()) {
                this.sessions.remove(key);
            }
        }
        LOGGER.debug("tftp session is end: " + session.isEnd() + ", current sessions count : " + this.sessions.size());
    }

    public static class T extends Thread {

        public T() {
            super("TTTTT");
        }

        @Override
        public void run() {
            while (true) {
            }
        }
    }

    /*public static void main(String[] args) {
        String dir = "D:\\cache\\";
        TFTPServer server = TFTPServer.of().mode(ServerMode.GET_AND_PUT).timeout(999999).readerFn((s)->{
            File file = new File(dir, s);
            if(!file.exists()){
                return null;
            }
            try {
                return new FileInputStream(file);
            } catch (FileNotFoundException e) {
                return null;
            }
        }).writerFn((s)->{
            FileOutputStream out;
            try {
                out = new FileOutputStream(new File(dir, s));
            } catch (FileNotFoundException e) {
                return null;
            }
            return out;
        });
        server.start();
    }*/
}
